import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import pandas as pd
from tqdm import tqdm

from preprocessing.option import parse_args, DEEPER_FORENSICS

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)


def get_real_fake_pairs(original_path, manipulated_path, args, mode):
    lines = open(os.path.join(original_path, '{}.txt'.format(mode)), 'r').readlines()
    video_names = [line[:-5] for line in lines]
    video_set = set(video_names)
    pairs = []

    for video_name in video_names:
        pair = (video_name, video_name)
        pairs.append(pair)

    for video_path in os.listdir(os.path.join(manipulated_path, args.fake_type)):
        video_path_split = video_path.split('_')
        assert len(video_path_split) == 2
        if video_path_split[0] not in video_set:
            continue
        pair = (video_path_split[0], video_path[:-4])
        pairs.append(pair)
    return pairs


def get_paths(vid, label, original_path, manipulated_path, args):
    ori_vid, fake_vid = vid
    ori_dir = os.path.join(original_path, 'crops', ori_vid)
    fake_dir = os.path.join(manipulated_path, args.crops_dir, fake_vid) if label == 1 \
        else os.path.join(original_path, 'crops', fake_vid)
    mask_dir = os.path.join(manipulated_path, args.diffs_dir, fake_vid)
    data = []
    # Some unaligned videos have been removed
    if label == 1 and not os.path.exists(fake_dir):
        return data
    # Some masks may have been removed
    if label == 1 and not os.path.exists(mask_dir):
        return data
    for frame in os.listdir(fake_dir):
        ori_img_path = os.path.join(ori_dir, frame)
        fake_img_path = os.path.join(fake_dir, frame)
        if label == 0 and os.path.exists(ori_img_path):
            data.append([ori_img_path, label, ori_vid])
        elif label == 1 and os.path.exists(fake_img_path) and os.path.exists(ori_img_path):
            mask_path = os.path.join(mask_dir, "{}_diff.png".format(frame[:-4]))
            if os.path.exists(mask_path):
                data.append([fake_img_path, label, ori_vid])

    return data


def save_folds(args, original_path, manipulated_path, ori_fake_pairs, mode='train'):
    data = []
    if mode == 'train_fake':
        ori_ori_pairs = set()
    else:
        ori_ori_pairs = set([(ori, ori) for ori, fake in ori_fake_pairs])
    ori_fake_pairs = set(ori_fake_pairs) - ori_ori_pairs
    with Pool(processes=1) as p:
        # original label=0
        with tqdm(total=len(ori_ori_pairs)) as pbar:
            func = partial(get_paths, label=0, original_path=original_path, manipulated_path=manipulated_path,
                           args=args)
            for v in p.imap_unordered(func, ori_ori_pairs):
                pbar.update()
                data.extend(v)
        # fake label=1
        with tqdm(total=len(ori_fake_pairs)) as pbar:
            func = partial(get_paths, label=1, original_path=original_path, manipulated_path=manipulated_path,
                           args=args)
            for v in p.imap_unordered(func, ori_fake_pairs):
                pbar.update()
                data.extend(v)
    fold_data = []
    for img_path, label, ori_vid in data:
        path = Path(img_path)
        video = path.parent.name
        file = path.name
        fold_data.append([video, file, label, ori_vid, int(file.split("_")[0])])
    # random.shuffle(fold_data)
    columns = ["video", "file", "label", "original", "frame"]
    save_file = '../data/{}/data_{}_{}_{}.csv'.format(DEEPER_FORENSICS, DEEPER_FORENSICS, args.fake_type, mode)
    pd.DataFrame(fold_data, columns=columns).to_csv(save_file, index=False)


def real_fake_split(ori_fake_pairs):
    ori_fake_pairs_real = []
    ori_fake_pairs_fake = []
    for pair in ori_fake_pairs:
        if pair[0] == pair[1]:
            ori_fake_pairs_real.append(pair)
        else:
            ori_fake_pairs_fake.append(pair)
    return ori_fake_pairs_real, ori_fake_pairs_fake


def main():
    args = parse_args()
    os.makedirs('../data/{}'.format(DEEPER_FORENSICS), exist_ok=True)
    original_path = os.path.join(args.ori_dir, 'original_sequences', 'youtube', 'c23')
    manipulated_path = os.path.join(args.root_dir, 'manipulated_videos')

    def cmp(pair):
        pair_fake_split = pair[1].split('_')
        if len(pair_fake_split) == 1:
            return pair[1]
        elif len(pair_fake_split) == 2:
            return '{}{}'.format(pair_fake_split[0], pair_fake_split[1])
        else:
            return '{}{}'.format(pair_fake_split[0], pair_fake_split[2])

    ori_fake_pairs_train = get_real_fake_pairs(original_path, manipulated_path, args, 'train')
    save_folds(args, original_path, manipulated_path, ori_fake_pairs_train, 'train')

    ori_fake_pairs_test = get_real_fake_pairs(original_path, manipulated_path, args, 'test')
    save_folds(args, original_path, manipulated_path, ori_fake_pairs_test, 'test')

    ori_fake_pairs_val = get_real_fake_pairs(original_path, manipulated_path, args, 'val')
    save_folds(args, original_path, manipulated_path, ori_fake_pairs_val, 'val')

    ori_fake_pairs_train.sort(key=cmp)
    ori_fake_pairs_real, ori_fake_pairs_fake = real_fake_split(ori_fake_pairs_train)
    save_folds(args, original_path, manipulated_path, ori_fake_pairs_real, 'train_real')
    save_folds(args, original_path, manipulated_path, ori_fake_pairs_fake, 'train_fake')


if __name__ == '__main__':
    # Total: 2000
    main()
